"""
Utilities for standardizing model outputs across the framework.

This module provides helper functions to ensure consistent output formats
across different model components, reducing code duplication and enforcing
interface contracts.
"""

from typing import Dict, Union, Any, Optional, TypeVar, cast, Type
import torch

from continuallearning.interfaces.types import (
    ModelOutput,
    BackboneOutput,
    AdapterOutput,
    HookOutput,
)

EMPTY_TENSOR = torch.zeros(0)

# 定义通用输出类型
OutputType = Union[
    torch.Tensor, Dict[str, Any], ModelOutput, AdapterOutput, BackboneOutput, HookOutput
]
T = TypeVar("T", ModelOutput, AdapterOutput, BackboneOutput, HookOutput)


def _extract_common_fields(
    output: OutputType,
) -> Dict[str, Any]:
    """
    提取各种输出格式中的共同字段。

    这是一个私有辅助函数，用于减少标准化函数之间的代码重复。

    Args:
        output: 原始输出(张量、字典或输出对象)

    Returns:
        Dict[str, Any]: 包含标准字段的字典
    """
    fields: Dict[str, Any] = {
        "features": EMPTY_TENSOR,
        "hidden_states": None,
        "attentions": None,
        "extra_outputs": None,
    }

    # 处理张量输入
    if isinstance(output, torch.Tensor):
        fields["features"] = output

    # 处理字典输入
    elif isinstance(output, dict):
        # 提取标准字段
        fields["features"] = output.get("features", EMPTY_TENSOR)
        fields["hidden_states"] = output.get("hidden_states")
        fields["attentions"] = output.get("attentions")

        # 尝试从last_hidden_state中提取特征
        if fields["features"].numel() == 0 and "last_hidden_state" in output:
            last_hidden_state = output["last_hidden_state"]
            if (
                isinstance(last_hidden_state, torch.Tensor)
                and last_hidden_state.numel() > 0
            ):
                fields["features"] = last_hidden_state[:, 0]  # 使用CLS token作为特征

        # 如果仍然没有特征，使用第一个可用的张量
        if fields["features"].numel() == 0:
            for key, value in output.items():
                if isinstance(value, torch.Tensor) and value.numel() > 0:
                    fields["features"] = value
                    break

    # 处理ModelOutput类型
    elif isinstance(output, ModelOutput):
        fields["features"] = (
            output.features if output.features.numel() > 0 else EMPTY_TENSOR
        )
        fields["hidden_states"] = output.hidden_states
        fields["attentions"] = output.attentions
        fields["extra_outputs"] = output.extra_outputs

    return fields


def standardize_backbone_output(
    output: Union[torch.Tensor, Dict[str, Any], ModelOutput, BackboneOutput],
) -> BackboneOutput:
    """
    将各种输出格式转换为标准化的BackboneOutput。

    Args:
        output: 骨干模型的原始输出(张量、字典或BackboneOutput)

    Returns:
        BackboneOutput: 具有一致属性的标准化输出
    """
    # 如果已经是正确的类型，直接返回
    if isinstance(output, BackboneOutput):
        return output

    # 提取通用字段
    fields = _extract_common_fields(output)

    # 提取骨干网络特有字段
    pooled_output = EMPTY_TENSOR
    last_hidden_state = EMPTY_TENSOR

    # 从字典或BackboneOutput中获取特殊字段
    if isinstance(output, dict):
        pooled_output = output.get("pooled_output", EMPTY_TENSOR)
        last_hidden_state = output.get("last_hidden_state", EMPTY_TENSOR)
    elif isinstance(output, BackboneOutput):
        pooled_output = output.pooled_output
        last_hidden_state = output.last_hidden_state

    return BackboneOutput(
        features=fields["features"],
        hidden_states=fields["hidden_states"],
        attentions=fields["attentions"],
        extra_outputs=fields["extra_outputs"],
        pooled_output=pooled_output,
        last_hidden_state=last_hidden_state,
    )


def standardize_adapter_output(
    output: Union[torch.Tensor, Dict[str, Any], ModelOutput, AdapterOutput],
) -> AdapterOutput:
    """
    将各种输出格式转换为标准化的AdapterOutput。

    Args:
        output: 适配器的原始输出(张量、字典或ModelOutput)

    Returns:
        AdapterOutput: 具有一致属性的标准化输出
    """
    # 如果已经是正确的类型，直接返回
    if isinstance(output, AdapterOutput):
        return output

    # 提取通用字段
    fields = _extract_common_fields(output)

    # 提取适配器特有字段
    adapter_hidden_states = None

    # 从字典或AdapterOutput中获取特殊字段
    if isinstance(output, dict):
        adapter_hidden_states = output.get("adapter_hidden_states")
    elif isinstance(output, AdapterOutput):
        adapter_hidden_states = output.adapter_hidden_states

    return AdapterOutput(
        features=fields["features"],
        hidden_states=fields["hidden_states"],
        attentions=fields["attentions"],
        extra_outputs=fields["extra_outputs"],
        adapter_hidden_states=adapter_hidden_states,
    )


def standardize_model_output(
    output: Union[torch.Tensor, Dict[str, Any], ModelOutput],
    logits: Optional[torch.Tensor] = None,
) -> ModelOutput:
    """
    将各种输出格式转换为标准化的ModelOutput。

    Args:
        output: 模型的原始输出(张量、字典或ModelOutput)
        logits: 可选的logits张量，包含在输出中

    Returns:
        ModelOutput: 具有一致属性的标准化输出
    """
    # 如果已经是正确类型且具有logits，直接返回
    if isinstance(output, ModelOutput) and (
        logits is None or (output.logits.numel() > 0)
    ):
        if logits is not None:
            # 创建带有提供的logits的新ModelOutput
            return ModelOutput(
                features=output.features,
                logits=logits,
                hidden_states=output.hidden_states,
                attentions=output.attentions,
                extra_outputs=output.extra_outputs,
            )
        return output

    # 提取通用字段
    fields = _extract_common_fields(output)

    # 处理logits
    out_logits = EMPTY_TENSOR if logits is None else logits

    # 从不同输入类型提取logits
    if isinstance(output, torch.Tensor) and logits is None:
        # 如果没有提供单独的logits，则将输出张量作为logits
        out_logits = output
    elif isinstance(output, dict) and logits is None:
        out_logits = output.get("logits", EMPTY_TENSOR)
    elif isinstance(output, ModelOutput):
        out_logits = logits if logits is not None else output.logits

    return ModelOutput(
        features=fields["features"],
        logits=out_logits,
        hidden_states=fields["hidden_states"],
        attentions=fields["attentions"],
        extra_outputs=fields["extra_outputs"],
    )


def standardize_hook_output(
    output: Union[Dict[str, Any], HookOutput],
) -> HookOutput:
    """
    将Hook输出格式标准化为HookOutput类型。

    Args:
        output: Hook的原始输出(通常是字典或HookOutput)

    Returns:
        HookOutput: 标准化的Hook输出
    """
    # 如果已经是正确的类型，直接返回
    if isinstance(output, HookOutput):
        return output

    # 处理字典输入
    if isinstance(output, dict):
        if "output" not in output:
            raise ValueError("Hook output dictionary must contain 'output' key")

        modified = output.get("modified", False)
        extra_data = {
            k: v for k, v in output.items() if k not in ("output", "modified")
        }
        extra_data = extra_data if extra_data else None

        return HookOutput(
            output=output["output"], modified=modified, extra_data=extra_data
        )

    raise TypeError(f"Cannot standardize hook output of type {type(output)}")


def standardize_output(
    output: OutputType,
) -> Union[ModelOutput, AdapterOutput, BackboneOutput, HookOutput]:
    """
    根据输入类型将各种输出格式转换为标准化输出。

    此函数作为分发器，根据输入类型和特征调用适当的专用标准化函数。

    Args:
        output: 要标准化的输出(张量、字典或数据类)

    Returns:
        Union[ModelOutput, AdapterOutput, BackboneOutput, HookOutput]: 标准化输出

    Raises:
        ValueError: 如果输出无法标准化
    """
    # 如果已经是标准化类型，直接返回
    if isinstance(output, (ModelOutput, AdapterOutput, BackboneOutput, HookOutput)):
        return output

    # 处理特殊的Hook输出
    if isinstance(output, dict) and "output" in output:
        return standardize_hook_output(output)

    # 处理字典类型，根据字典内容决定使用哪个专用函数
    if isinstance(output, dict):
        # 检查是否有BackboneOutput特有字段
        if "pooled_output" in output or "last_hidden_state" in output:
            return standardize_backbone_output(output)
        # 检查是否有AdapterOutput特有字段
        elif "adapter_hidden_states" in output:
            return standardize_adapter_output(output)
        # 检查是否有ModelOutput特有字段
        elif "logits" in output:
            return standardize_model_output(output)

    # 默认情况下，对于tensor和其他类型，使用adapter输出标准化
    # 这保持了与之前代码的兼容性
    return standardize_adapter_output(output)


def convert_output_type(output: OutputType, target_class: Type[T]) -> T:
    """
    将输出转换为指定的目标类型。

    这个函数允许在不同的输出类型之间进行显式转换，
    在需要特定类型输出但实际收到另一种类型时非常有用。

    Args:
        output: 输入的输出对象
        target_class: 目标输出类的类型 (ModelOutput, AdapterOutput等类，不是函数)

    Returns:
        T: 转换后的输出对象

    Raises:
        TypeError: 如果转换不可行
    """
    # 如果已经是目标类型，直接返回
    if isinstance(output, target_class):
        return cast(T, output)

    # 先标准化输出
    std_output = standardize_output(output)

    # FIXME Any 还是尽量少用，想想办法
    fields: Dict[str, Any] = {
        "features": getattr(std_output, "features", EMPTY_TENSOR),
        "hidden_states": getattr(std_output, "hidden_states", None),
        "attentions": getattr(std_output, "attentions", None),
        "extra_outputs": getattr(std_output, "extra_outputs", None),
    }

    # 添加特定类型的字段
    if target_class is BackboneOutput:
        fields["pooled_output"] = (
            getattr(std_output, "pooled_output", EMPTY_TENSOR)
            if hasattr(std_output, "pooled_output")
            else EMPTY_TENSOR
        )
        fields["last_hidden_state"] = (
            getattr(std_output, "last_hidden_state", EMPTY_TENSOR)
            if hasattr(std_output, "last_hidden_state")
            else EMPTY_TENSOR
        )
    elif target_class is AdapterOutput:
        fields["adapter_hidden_states"] = (
            getattr(std_output, "adapter_hidden_states", None)
            if hasattr(std_output, "adapter_hidden_states")
            else None
        )
    elif target_class is ModelOutput:
        fields["logits"] = (
            getattr(std_output, "logits", EMPTY_TENSOR)
            if hasattr(std_output, "logits")
            else EMPTY_TENSOR
        )
    elif target_class is HookOutput:
        # HookOutput是完全不同的结构，需要特殊处理
        if isinstance(std_output, HookOutput):
            return cast(T, std_output)
        raise TypeError(f"Cannot convert {type(std_output)} to HookOutput")

    # 创建目标类型的实例
    return target_class(**fields)
